import numpy as np
from pysr import PySRRegressor
from bishe_situations.utils import generate_data, plot_results, evaluate_model

# 定义目标函数
def target_function(X):
    return np.sin(X) + X**2

# 生成数据
X, y = generate_data(target_function, n_samples=1000, noise=0.1)
unaries = ["neg", "exp", "sin", "sinh", "erf", "round",
           "square", "log", "cos", "cosh", "erfc", "floor",
           "cube", "log10", "tan", "tanh", "gamma", "ceil", 
           "cbrt", "log2", "asin", "asinh", "relu",
           "sqrt", "log1p", "acos", "acosh", "sinc",
           "abs", "atan", "atanh",
           "sign",
           "inv"]

# 创建PySR模型，使用最简单的配置
model = PySRRegressor(
    binary_operators=["+", "-", "*", "/", "^"],
    unary_operators=unaries,
    niterations=20,
    population_size=20,
    maxsize=20,
    parsimony=0.1,
    verbosity=1
)

# 训练模型
model.fit(X, y)

# 预测
y_pred = model.predict(X)

# 评估模型
mse, r2 = evaluate_model(y, y_pred)

# 打印找到的表达式
print("\n找到的表达式:")
print(model.sympy())

# 绘制结果
plot_results(X, y, y_pred, "PySR符号回归结果") 